import tushare as ts
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt

# 配置中文字体和负号显示
plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False

class DataFetcher:
    def __init__(self, token):
        ts.set_token(token)
        self.pro = ts.pro_api()

    def get_price_data(self, ts_codes, start_date, end_date):
        all_data = {}
        for code in ts_codes:
            df = self.pro.daily(ts_code=code, start_date=start_date, end_date=end_date)
            if df.empty:
                continue
            df = df.sort_values('trade_date')
            df['trade_date'] = pd.to_datetime(df['trade_date'])
            df.set_index('trade_date', inplace=True)
            all_data[code] = df
        return all_data

class Strategy:
    def __init__(self):
        self.short_window = 5
        self.long_window = 20
        self.rsi_period = 14
        self.rsi_oversold = 20
        self.rsi_overbought = 80

    def prepare(self, df):
        df = df.copy()
        df['short_ma'] = df['close'].rolling(window=self.short_window).mean()
        df['long_ma'] = df['close'].rolling(window=self.long_window).mean()
        delta = df['close'].diff()
        gain = delta.where(delta > 0, 0)
        loss = -delta.where(delta < 0, 0)
        avg_gain = gain.rolling(self.rsi_period).mean()
        avg_loss = loss.rolling(self.rsi_period).mean()
        rs = avg_gain / avg_loss
        df['rsi'] = 100 - 100 / (1 + rs)
        return df

    def generate_signal(self, df):
        df = self.prepare(df)
        if len(df) < self.long_window + 1:
            return 0  # 数据不够，不能操作

        # 取最近两天数据判断金叉死叉
        yesterday = df.iloc[-2]
        today = df.iloc[-1]

        ma_signal = 0
        if yesterday['short_ma'] < yesterday['long_ma'] and today['short_ma'] > today['long_ma']:
            ma_signal = 1  # 金叉买入信号
        elif yesterday['short_ma'] > yesterday['long_ma'] and today['short_ma'] < today['long_ma']:
            ma_signal = -1  # 死叉卖出信号

        rsi_signal = 0
        if today['rsi'] < self.rsi_oversold:
            rsi_signal = 1  # 超卖买入
        elif today['rsi'] > self.rsi_overbought:
            rsi_signal = -1  # 超买卖出

        # 综合信号
        if ma_signal == 1 and rsi_signal == 1:
            return 2  # 强买入
        elif ma_signal == -1 and rsi_signal == -1:
            return -2  # 强卖出
        elif ma_signal == 1 or rsi_signal == 1:
            return 1   # 买入
        elif ma_signal == -1 or rsi_signal == -1:
            return -1  # 卖出
        else:
            return 0  # 无操作

class RiskManager:
    def __init__(self, stop_loss=0.05, take_profit=0.10):
        self.stop_loss = stop_loss
        self.take_profit = take_profit

    def check(self, entry_price, current_price):
        ret = (current_price - entry_price) / entry_price
        if ret <= -self.stop_loss:
            return 'stop_loss'
        elif ret >= self.take_profit:
            return 'take_profit'
        else:
            return 'hold'

class Portfolio:
    def __init__(self, initial_cash=1000000):
        self.cash = initial_cash
        self.positions = {}  # key: ts_code, value: dict with shares, entry_price, entry_date
        self.history = []
        self.net_values = []

    def update(self, prices, current_date):
        total_value = self.cash
        for code, pos in self.positions.items():
            total_value += pos['shares'] * prices.get(code, 0)
        self.net_values.append(total_value)

    def trade(self, code, signal, price, date):
        if signal > 0:
            if code in self.positions:
                return  # 已持仓，暂不加仓
            shares = int(self.cash // price)
            if shares > 0:
                self.positions[code] = {'shares': shares, 'entry_price': price, 'entry_date': date}
                self.cash -= shares * price
                self.history.append({'date': date, 'code': code, 'action': 'buy', 'price': price, 'shares': shares})
        elif signal < 0:
            if code in self.positions:
                pos = self.positions.pop(code)
                self.cash += pos['shares'] * price
                self.history.append({'date': date, 'code': code, 'action': 'sell', 'price': price, 'shares': pos['shares'],
                                     'holding_days': (date - pos['entry_date']).days})

class Backtest:
    def __init__(self, token, stocks, start_date, end_date):
        self.data_fetcher = DataFetcher(token)
        self.stocks = stocks
        self.start_date = start_date
        self.end_date = end_date
        self.strategy = Strategy()
        self.risk_manager = RiskManager()
        self.portfolio = Portfolio()

    def run(self):
        price_data = self.data_fetcher.get_price_data(self.stocks, self.start_date, self.end_date)
        all_dates = sorted({date for df in price_data.values() for date in df.index})

        for date in all_dates:
            current_prices = {}
            signals = {}

            for code in self.stocks:
                if code not in price_data or date not in price_data[code].index:
                    continue
                df = price_data[code][price_data[code].index <= date]
                current_price = price_data[code].loc[date, 'close']
                current_prices[code] = current_price

                signal = self.strategy.generate_signal(df)

                # 止损止盈优先级最高
                if code in self.portfolio.positions:
                    pos = self.portfolio.positions[code]
                    risk_status = self.risk_manager.check(pos['entry_price'], current_price)
                    if risk_status == 'stop_loss' or risk_status == 'take_profit':
                        signal = -2  # 强制卖出

                signals[code] = signal

            # 执行交易
            for code, signal in signals.items():
                if signal != 0:
                    self.portfolio.trade(code, signal, current_prices[code], date)

            self.portfolio.update(current_prices, date)

        return self._summary()

    def _summary(self):
        start_value = self.portfolio.net_values[0]
        end_value = self.portfolio.net_values[-1]
        total_return = (end_value - start_value) / start_value * 100
        print(f"初始资金: {start_value:.2f} 元")
        print(f"最终资金: {end_value:.2f} 元")
        print(f"总收益率: {total_return:.2f} %")
        print(f"交易次数: {len(self.portfolio.history)}")
        return {
            'start_value': start_value,
            'end_value': end_value,
            'total_return_percent': total_return,
            'trade_count': len(self.portfolio.history),
            'net_values': self.portfolio.net_values
        }

def plot_results(net_values):
    plt.figure(figsize=(12,6))
    plt.plot(net_values, label='组合净值')
    plt.title('回测净值曲线')
    plt.xlabel('交易天数')
    plt.ylabel('资金净值')
    plt.legend()
    plt.grid(True)
    plt.show()


def plot_net_value(dates, net_values, code=None):
    plt.figure(figsize=(12, 6))
    plt.plot(dates, net_values, label='净值曲线')
    plt.xlabel('日期')
    plt.ylabel('组合净值')
    plt.title(f'策略净值曲线 {f"- {code}" if code else ""}')
    plt.legend()
    plt.grid(True)
    plt.show()



def plot_price_signals(df, buy_dates, sell_dates, code):
    plt.figure(figsize=(14, 7))
    plt.plot(df.index, df['close'], label=f'{code} 收盘价')
    plt.scatter(buy_dates, df.loc[buy_dates, 'close'], marker='^', color='green', label='买入', s=100)
    plt.scatter(sell_dates, df.loc[sell_dates, 'close'], marker='v', color='red', label='卖出', s=100)
    plt.title(f'{code} 收盘价及交易信号')
    plt.xlabel('日期')
    plt.ylabel('价格')
    plt.legend()
    plt.grid()
    plt.show()



def plot_trade_returns(trade_returns):
    plt.figure(figsize=(10, 5))
    plt.hist(trade_returns, bins=30, alpha=0.7, color='blue')
    plt.xlabel('单笔交易收益率')
    plt.ylabel('频数')
    plt.title('单笔交易收益分布图')
    plt.grid(True)
    plt.show()


def plot_multiple_stock_returns(stock_returns):
    codes = list(stock_returns.keys())
    returns = [stock_returns[c] for c in codes]
    plt.figure(figsize=(12, 6))
    plt.bar(codes, returns, color='skyblue')
    plt.xlabel('股票代码')
    plt.ylabel('总收益率（%）')
    plt.title('多只股票回测收益对比')
    plt.xticks(rotation=45)
    plt.grid(axis='y')
    plt.show()



# ==== 策略实现函数 ====
def run_strategy(code):
    df = ts.pro_bar(ts_code=code, start_date='20220101', end_date='20250101', adj='qfq')
    df = df.sort_values('trade_date')
    df.set_index(pd.to_datetime(df['trade_date']), inplace=True)

    df['ma10'] = df['close'].rolling(10).mean()
    df['ma30'] = df['close'].rolling(30).mean()

    df['signal'] = 0
    df.loc[df['ma10'] > df['ma30'], 'signal'] = 1
    df.loc[df['ma10'] < df['ma30'], 'signal'] = -1
    df['position'] = df['signal'].shift(1)

    df['daily_ret'] = df['close'].pct_change()
    df['strategy_ret'] = df['daily_ret'] * df['position']
    df['net_value'] = (1 + df['strategy_ret']).cumprod()

    # 提取买卖点
    buy_dates = df[(df['position'] == 1) & (df['position'].shift(1) != 1)].index
    sell_dates = df[(df['position'] == -1) & (df['position'].shift(1) != -1)].index

    # 单笔收益列表
    trade_returns = []
    for b, s in zip(buy_dates, sell_dates):
        if s > b:
            ret = df.loc[s, 'close'] / df.loc[b, 'close'] - 1
            trade_returns.append(ret)

    # ===== 调用图表函数展示结果 =====
    print(f'\n📊 股票 {code} 策略收益率：{(df["net_value"].iloc[-1] - 1) * 100:.2f}%')
    plot_net_value(df.index, df['net_value'], code)
    plot_price_signals(df, buy_dates, sell_dates, code)
    if trade_returns:
        plot_trade_returns(trade_returns)

    return (df["net_value"].iloc[-1] - 1) * 100  # 返回收益率


# ==== 多股票批量回测 ====
stock_list = ['000001.SZ', '000002.SZ', '000651.SZ']  # 示例股票
stock_returns = {}

for code in stock_list:
    try:
        ret = run_strategy(code)
        stock_returns[code] = ret
    except Exception as e:
        print(f'跳过 {code}，错误：{e}')

# 所有股票的对比图
plot_multiple_stock_returns(stock_returns)
if __name__ == '__main__':
    TOKEN = '1c7f85b9026518588c0d0cdac712c2d17344332c9c8cfe6bc83ee75c'
    stocks = ['000001.SZ', '600519.SH', '000651.SZ']  # 你想回测的股票列表
    start_date = '20220101'
    end_date = '20231231'

    backtest = Backtest(TOKEN, stocks, start_date, end_date)
    results = backtest.run()

    if results['total_return_percent'] > 0:
        print("策略收益为正，回测成功。")
    else:
        print("策略收益未达预期，请优化策略。")
